# -*- encoding: utf-8 -*-
"""
@author: acedar  
@time: 2023/10/5 13:04
@file: main.py
"""

import numpy as np
import torch
from torch.autograd import Variable
from modeling import subsequent_mask, loss_backprop, make_std_mask, make_model, get_std_opt, LabelSmoothing
np.random.seed(0)


def train_epoch(train_iter, model, criterion, opt, transpose=False):
    model.train()
    for i, batch in enumerate(train_iter):
        src, trg, src_mask, trg_mask = \
            batch.src, batch.trg, batch.src_mask, batch.trg_mask

        # src:
        # tensor([[ 6,  1,  4,  4,  8],
        #         [10,  4,  6,  3,  5],
        #         [ 8,  7,  9,  9,  2],
        #         [ 7,  8,  8,  9,  2],
        #         [ 6, 10,  9, 10,  5],
        #         [ 4,  1,  4,  6,  1],
        #         [ 3,  4,  9,  2,  4],
        #         [ 4,  4,  8,  1,  2]], dtype=torch.int32)

        # trg:
        # tensor([[ 6,  1,  4,  4,  8],
        #         [10,  4,  6,  3,  5],
        #         [ 8,  7,  9,  9,  2],
        #         [ 7,  8,  8,  9,  2],
        #         [ 6, 10,  9, 10,  5],
        #         [ 4,  1,  4,  6,  1],
        #         [ 3,  4,  9,  2,  4],
        #         [ 4,  4,  8,  1,  2]], dtype=torch.int32)

        # src_mask:
        # tensor([[[True, True, True, True, True]],
        #         [[True, True, True, True, True]],
        #         [[True, True, True, True, True]],
        #         [[True, True, True, True, True]],
        #         [[True, True, True, True, True]],
        #         [[True, True, True, True, True]],
        #         [[True, True, True, True, True]],
        #         [[True, True, True, True, True]]])

        # trg_mask:
        # tensor([[[ True, False, False, False, False],
        #          [ True,  True, False, False, False],
        #          [ True,  True,  True, False, False],
        #          [ True,  True,  True,  True, False],
        #          [ True,  True,  True,  True,  True]],
        #
        #         [[ True, False, False, False, False],
        #          [ True,  True, False, False, False],
        #          [ True,  True,  True, False, False],
        #          [ True,  True,  True,  True, False],
        #          [ True,  True,  True,  True,  True]],
        #
        #         [[ True, False, False, False, False],
        #          [ True,  True, False, False, False],
        #          [ True,  True,  True, False, False],
        #          [ True,  True,  True,  True, False],
        #          [ True,  True,  True,  True,  True]],
        #
        #         [[ True, False, False, False, False],
        #          [ True,  True, False, False, False],
        #          [ True,  True,  True, False, False],
        #          [ True,  True,  True,  True, False],
        #          [ True,  True,  True,  True,  True]],
        #
        #         [[ True, False, False, False, False],
        #          [ True,  True, False, False, False],
        #          [ True,  True,  True, False, False],
        #          [ True,  True,  True,  True, False],
        #          [ True,  True,  True,  True,  True]],
        #
        #         [[ True, False, False, False, False],
        #          [ True,  True, False, False, False],
        #          [ True,  True,  True, False, False],
        #          [ True,  True,  True,  True, False],
        #          [ True,  True,  True,  True,  True]],
        #
        #         [[ True, False, False, False, False],
        #          [ True,  True, False, False, False],
        #          [ True,  True,  True, False, False],
        #          [ True,  True,  True,  True, False],
        #          [ True,  True,  True,  True,  True]],
        #
        #         [[ True, False, False, False, False],
        #          [ True,  True, False, False, False],
        #          [ True,  True,  True, False, False],
        #          [ True,  True,  True,  True, False],
        #          [ True,  True,  True,  True,  True]]])

        # forward.input: torch.Size([16, 5]) torch.Size([16, 4]) torch.Size([16, 1, 5]) torch.Size([16, 4, 4])
        print("forward.input:", src.size(), trg[:, :-1].size(), src_mask.size(), trg_mask[:, :-1, :-1].size())
        # print("trg[:, :-1]", trg[:, :-1])
        # print("trg_mask[:, :-1, :-1]", trg_mask[:, :-1, :-1])

        out = model.forward(src, trg[:, :-1], src_mask, trg_mask[:, :-1, :-1])
        # out = model.forward(src, trg, src_mask, trg_mask)

        # out: [16, 4, 512]
        print("out:", out.size())
        loss = loss_backprop(model.generator, criterion, out, trg[:, 1:], batch.ntokens)
        # loss = loss_backprop(model.generator, criterion, out, trg, batch.ntokens)

        model_opt.step()
        model_opt.optimizer.zero_grad()
        if i % 10 == 1:
            print(i, loss, model_opt._rate)


def valid_epoch(valid_iter, model, criterion, transpose=False):
    model.test()
    total = 0
    for batch in valid_iter:
        src, trg, src_mask, trg_mask = \
            batch.src, batch.trg, batch.src_mask, batch.trg_mask
        out = model.forward(src, trg[:, :-1], src_mask, trg_mask[:, :-1, :-1])
        loss = loss_backprop(model.generator, criterion, out, trg[:, 1:], batch.ntokens)


class Batch:
    def __init__(self, src, trg, src_mask, trg_mask, ntokens):
        self.src = src
        self.trg = trg
        self.src_mask = src_mask
        self.trg_mask = trg_mask
        self.ntokens = ntokens


def data_gen(V, batch, nbatches):
    for i in range(nbatches):
        data = torch.from_numpy(np.random.randint(1, V, size=(batch, 5)))
        src = Variable(data, requires_grad=False)
        tgt = Variable(data, requires_grad=False)
        # print("src:", src.size(), src.dim(), src)
        # print("tgt:", tgt.size(), tgt.dim(), tgt)

        src_mask, tgt_mask = make_std_mask(src, tgt, 0)
        yield Batch(src, tgt, src_mask, tgt_mask, (tgt[1:] != 0).data.sum())


V = 11
criterion = LabelSmoothing(size=V, padding_idx=0, smoothing=0.0)
model = make_model(V, V, N=1)
model_opt = get_std_opt(model)
for epoch in range(1):
    train_epoch(data_gen(V, 8, 1), model, criterion, model_opt)
